# Title     : TODO
# Objective : TODO
# Created by: Administrator
# Created on: 2019/7/24

library(ggrepel)
library(ropls)
library(pROC)
library(egg)
library(randomForest)
library(Boruta)
library(magrittr)
library(optparse)
library(gbm)
library(caret)
library(tidyverse)
library(readxl)
library(writexl)

createWhenNoExist <- function(f) {
  !dir.exists(f) && dir.create(f)
}

option_list <- list(
  make_option("--i", default = "../../data.xlsx", type = "character", help = "metabolite data file"),
  make_option("--g", default = "../../group.xlsx", type = "character", help = "sample group file"),
  make_option("--sc", default = "sample_color.txt", type = "character", help = "sample color file")
)
opt <- parse_args(OptionParser(option_list = option_list))

options(digits = 3)

sampleInfo <- read_xlsx(opt$g) %>%
  rename(ClassNote = Group) %>%
  select(c("SampleID", "ClassNote"))

head(sampleInfo)

parent <- paste0("./")
createWhenNoExist(parent)

numColumns <- c("年龄", "BMI", "腰围", "阿司匹林服药频率为每周几次", "您从多少岁开始养成每天都吸烟的习惯？",
                "扣除戒烟年数，您一共吸烟多少年？", "您平均每天吸烟多少支？",
                "在有烟雾的室内环境中，您与吸烟者共同居住或工作了多少年", "扣除戒烟年数，该吸烟者一共吸烟多少年",
                "该吸烟者平均每天吸烟多少支", "您每周有几天饮酒？", "白酒每周饮酒量", "您每周参加体育锻炼累计多长时间？")

originalData <- read_xlsx(opt$i) %>%
  rename(Metabolite = SampleID) %>%
  gather("SampleID", "Value", -Metabolite) %>%
  spread(Metabolite, "Value") %>%
  inner_join(sampleInfo, by = c("SampleID"))

originalColumnNames <- colnames(originalData)

data <- originalData %>%
  mutate(ClassNote = factor(ClassNote, levels = unique(ClassNote)))
allTrainData <- data %>%
  column_to_rownames("SampleID")

load("../../folds.RData")

predictLists <- (1:k) %>%
  map(function(k) {
    testData <- allTrainData[flds[[k]],] %>%
      mutate_at(vars(numColumns), as.numeric) %>%
      mutate_at(vars(-numColumns), factor)
    trainData <- allTrainData[-flds[[k]],] %>%
      mutate_at(vars(numColumns), as.numeric) %>%
      mutate_at(vars(-numColumns), factor)

    filterColumns <- trainData %>%
      sapply(nlevels) %>%
      keep(function(x) { x == 1 }) %>%
      names()
    print(filterColumns)

    trainData <- trainData %>%
      select(-filterColumns)

    testData <- testData %>%
      select(-filterColumns)

    common <- intersect(names(testData), names(trainData))
    for (p in common) {
      if (class(trainData[[p]]) == "factor") {
        levels(testData[[p]]) <- levels(trainData[[p]])
      }
    }

    glmRs <- glm(ClassNote ~ ., data = trainData, family = binomial(link = "logit"))
    Yhat <- predict(glmRs, newdata = testData, type = "response")

    YDf <- Yhat %>%
      as.data.frame() %>%
      rownames_to_column("SampleID") %>%
      set_colnames(c("SampleID", "Value"))

    thresh <- 0.5
    uniq.group <- unique(originalData$ClassNote)
    Yfac <- as.factor(testData$ClassNote)
    print("===")
    print(head(uniq.group))
    YhatFac <- cut(Yhat, breaks = c(-Inf, thresh, Inf), labels = c(uniq.group[1], uniq.group[2]))
    sum <- length(YhatFac)
    count <- 0

    for (i in 1:length(YhatFac)) {
      if (as.character(YhatFac[i]) == as.character(Yfac[i])) {
        count = count + 1
      }
    }
    if (count / sum < 0.5) {
      YhatFac = cut(Yhat, breaks = c(-Inf, thresh, Inf), labels = c(uniq.group[2], uniq.group[1]))
    }

    predictDf1 <- YDf %>%
      mutate(LR_prediction = YhatFac)
    predictFinalDf1 <- sampleInfo %>%
      left_join(predictDf1, by = c("SampleID")) %>%
      mutate(Fold = k)

    predictDf <- YDf %>%
      mutate(Prediction = YhatFac)
    predictFinalDf <- sampleInfo %>%
      left_join(predictDf, by = c("SampleID")) %>%
      rename(Sample = SampleID) %>%
      rename(Probability = Value) %>%
      mutate_at(vars("Probability"), function(x) {
        ifelse(x > 0.5, x, 1 - x)
      }) %>%
      select(-c("Probability"), "Probability") %>%
      mutate(Test_in_CV_fold = k)

    list(Yfac = Yfac, YhatFac = YhatFac, plotData = predictFinalDf1, predictDf = predictFinalDf)
  })

Yfac <- predictLists %>%
  map(function(list) {
    list$Yfac %>%
      as.character()
  }) %>%
  flatten_chr()

YhatFac <- predictLists %>%
  map(function(list) {
    list$YhatFac %>%
      as.character()
  }) %>%
  flatten_chr()

pre_summary = table(YhatFac, Yfac)

summaryTb <- pre_summary %>%
  as.data.frame() %>%
  as_tibble() %>%
  rename(Var1 = 1, Var2 = 2) %>%
  spread(Var2, "Freq")
print("=log=")
print(summaryTb)
summaryMatrix <- summaryTb %>%
  select(-"Var1") %>%
  as.matrix()
diagSum <- sum(diag(summaryMatrix))
sum <- sum(summaryMatrix)
predictive <- (diagSum / sum) %>%
  round(3)
finalSummaryTb <- summaryTb %>%
  mutate(`Model predictive accuracy` = c(predictive, "")) %>%
  rename(` ` = Var1)
write_csv(finalSummaryTb, "LR_Prediction_Summary_CV_Overall.csv")

predictDf <- predictLists %>%
  map_dfr(function(list) {
    list$predictDf
  })
predictDf
write.csv(predictDf, "LR_Prediction_CV.csv", row.names = F)

plotData <- predictLists %>%
  map_dfr(function(list) {
    list$plotData
  })
plotData

write_tsv(plotData, "Classification_Result.txt")











